import asyncio
import requests
import logging
from typing import Optional
from dataclasses import dataclass
import time
import tiktoken  # 新增token计算库
from utils.logger import logger
import json

@dataclass
class APIResponse:
    success: bool
    content: Optional[str] = None
    error: Optional[str] = None

class APIClient:
    def __init__(self, config: dict):
        self.config = config
        # 确保API密钥正确
        if not config['openrouter'].get('api_key'):
            raise ValueError("缺少OpenRouter API密钥，请检查配置文件")
        
        # 添加必要的HTTP头
        self.headers = {
            "Authorization": f"Bearer {config['openrouter']['api_key']}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://your-domain.com",  # OpenRouter要求
            "X-Title": "Novel Rewriter"  # 应用标识
        }
        self.base_url = config['openrouter']['base_url']
        self.timeout = 30
        self.retry_config = {
            'max_retries': config['processing']['max_retries'],
            'retry_delay': 2,
            'retry_codes': {429, 502, 503}
        }
        self.encoder = tiktoken.get_encoding("cl100k_base")  # 与GPT-4相同的编码器

    def rewrite(self, text: str) -> str:
        logger.debug(f"开始处理文本块 | 长度: {len(text)}字符 | 前50字: {text[:50]}...")
        logger.trace(f"完整分块内容:\n{text}")  # 记录完整分块内容
        
        # 在发送请求前进行token校验
        input_tokens = len(self.encoder.encode(text))
        logger.debug(f"估算Token使用 | 输入: {input_tokens} | 最大输出: {self._calculate_max_tokens(text)}")
        
        # 记录完整请求载荷
        payload = self._build_payload(text)
        logger.trace("完整请求载荷:\n" + json.dumps(payload, indent=2, ensure_ascii=False))
        
        for attempt in range(1, self.retry_config['max_retries']+1):
            try:
                response = requests.post(
                    self.base_url,
                    headers=self.headers,
                    json=payload,
                    timeout=self.timeout
                )
                logger.debug(f"响应状态码: {response.status_code} | 响应长度: {len(response.text)}字节")
                logger.trace(f"完整响应内容:\n{response.text}")  # 记录完整响应
                response.raise_for_status()
                result = response.json()['choices'][0]['message']['content']
                
                # 新增结果校验
                if len(result) < len(text) * 0.3:  # 结果过短
                    raise ValueError("返回内容过短")
                if "..." in result[-3:]:  # 截断标记
                    raise ValueError("检测到截断内容")
                    
                if response.json().get('usage'):
                    logging.debug(f"Token使用: 输入{response.json()['usage']['prompt_tokens']} 输出{response.json()['usage']['completion_tokens']}")
                
                # 增强响应解析
                if not isinstance(result, str):
                    raise ValueError("返回内容格式错误")
                    
                # 检查是否包含原文关键词
                keywords = ["小明", "书店"]  # 根据输入动态生成
                if not any(kw in result for kw in keywords):
                    logger.warning("返回内容可能未保留关键信息")
                
                logger.success(f"处理完成 | 结果长度: {len(result)}字符")
                return result
            except Exception as e:
                self._handle_retry(attempt, str(e))
        return text  # 最终失败返回原文

    def _calculate_max_tokens(self, text: str) -> int:
        """精确计算可用token数"""
        # 计算输入token数（1个中文字符≈1.5token）
        input_tokens = len(self.encoder.encode(text))
        
        # 保留至少20%的token给输出
        max_output = min(
            int((self.config['openrouter']['max_tokens'] - input_tokens) * 0.8),
            int(input_tokens * self.config['processing']['token_ratio'])
        )
        return max(1024, max_output)  # 提高最小输出长度

    def _build_payload(self, text: str) -> dict:
        """构造API请求载荷"""
        return {
            "model": self.config['openrouter']['model'],
            "messages": [{
                "role": "user",
                "content": (
                    "你是一个专业的小说编辑，请严格遵循以下要求处理文本：\n"
                    "1. 必须保留所有人物名称、关键地点和核心情节\n"
                    "2. 优化语言表达但不要删减内容\n"
                    "3. 保持原有段落结构（用两个换行符分隔段落）\n"
                    "4. 输出长度必须与原文相当（±20%）\n"
                    "需要处理的原文：\n"
                    "-----\n"
                    f"{text}\n"
                    "-----\n"
                    "请返回优化后的版本（直接输出内容，不要添加任何解释）："
                )
            }],
            "temperature": self.config['processing']['temperature'],
            "max_tokens": self._calculate_max_tokens(text)
        }

    def _handle_rate_limit(self, response):
        retry_after = int(response.headers.get('Retry-After', 5))
        time.sleep(retry_after)
        return APIResponse(False, error=f"Rate limited, retry after {retry_after}s")

    def _handle_retry(self, attempt: int, error_msg: str):
        delay = self.retry_config['retry_delay'] ** attempt
        logging.warning(f"第{attempt}次重试失败: {error_msg}, {delay}秒后重试")
        time.sleep(delay)

    def _handle_response(self, response, original_text) -> APIResponse:
        """处理API响应"""
        start_time = time.time()
        try:
            if response.status == 200:
                data = response.json()
                logging.debug(f"API请求成功 | 耗时 {time.time()-start_time:.2f}s | 响应长度：{len(data)} 字节")
                return APIResponse(
                    success=True,
                    content=data['choices'][0]['message']['content'].strip()
                )
                
            error_info = f"状态码: {response.status} 响应: {response.text()}"
            if response.status in self.retry_config['retry_codes']:
                return self._handle_rate_limit(response)
                
            logging.error(f"不可恢复错误: {error_info}")
            return APIResponse(False, error=error_info, content=original_text)
            
        except (KeyError, ValueError) as e:
            error_msg = f"响应解析失败: {str(e)}"
            logging.error(error_msg)
            return APIResponse(False, error=error_msg, content=original_text)

    def debug_rewrite(self, text: str):
        """调试模式：打印完整请求信息"""
        payload = self._build_payload(text)
        print("=== 调试模式 ===")
        print(f"请求URL: {self.base_url}")
        print(f"请求头: {self.headers}")
        print(f"请求体: {json.dumps(payload, indent=2, ensure_ascii=False)}")
        
        # 发送测试请求
        response = requests.post(self.base_url, headers=self.headers, json=payload)
        print(f"响应状态码: {response.status_code}")
        print(f"响应内容: {response.text}")
        
        # 添加响应内容分析
        try:
            result = response.json()['choices'][0]['message']['content']
            print("\n响应分析：")
            print(f"输入长度: {len(text)} 字符")
            print(f"输出长度: {len(result)} 字符")
            print(f"关键内容保留: {'是' if '小明' in result else '否'}")
        except:
            print("无法解析响应内容")
